/*
 * Copyright 2002-2018 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package cn.cloud.all.security.oauth2.client.authentication;

import cn.cloud.all.security.authentication.ReactiveAuthenticationManager;
import cn.cloud.all.security.core.Authentication;
import cn.cloud.all.security.core.OAuth2AccessToken;
import cn.cloud.all.security.core.OAuth2RefreshToken;
import cn.cloud.all.security.core.endpoint.OAuth2AccessTokenResponse;
import cn.cloud.all.security.core.endpoint.OAuth2AuthorizationExchange;
import cn.cloud.all.security.oauth2.endpoint.OAuth2AuthorizationCodeGrantRequest;
import cn.cloud.all.security.oauth2.endpoint.ReactiveOAuth2AccessTokenResponseClient;
import org.springframework.util.Assert;
import reactor.core.publisher.Mono;

import java.util.function.Function;

public class OAuth2AuthorizationCodeReactiveAuthenticationManager implements ReactiveAuthenticationManager {
    private final ReactiveOAuth2AccessTokenResponseClient<OAuth2AuthorizationCodeGrantRequest> accessTokenResponseClient;

    public OAuth2AuthorizationCodeReactiveAuthenticationManager(ReactiveOAuth2AccessTokenResponseClient<OAuth2AuthorizationCodeGrantRequest> accessTokenResponseClient) {
        Assert.notNull(accessTokenResponseClient, "accessTokenResponseClient cannot be null");
        this.accessTokenResponseClient = accessTokenResponseClient;
    }

    @Override
    public Mono<Authentication> authenticate(Authentication authentication) {
        return Mono.defer(() -> {
            OAuth2AuthorizationCodeAuthenticationToken token = (OAuth2AuthorizationCodeAuthenticationToken) authentication;

            OAuth2AuthorizationExchangeValidator.validate(token.getAuthorizationExchange());

            OAuth2AuthorizationCodeGrantRequest authzRequest = new OAuth2AuthorizationCodeGrantRequest(token.getClientRegistration(), token.getAuthorizationExchange());

            return this.accessTokenResponseClient.getTokenResponse(authzRequest).map(onSuccess(token));
        });
    }

    private Function<OAuth2AccessTokenResponse, OAuth2AuthorizationCodeAuthenticationToken> onSuccess(OAuth2AuthorizationCodeAuthenticationToken token) {
        return accessTokenResponse -> {
            ClientRegistration registration = token.getClientRegistration();
            OAuth2AuthorizationExchange exchange = token.getAuthorizationExchange();
            OAuth2AccessToken accessToken = accessTokenResponse.getAccessToken();
            OAuth2RefreshToken refreshToken = accessTokenResponse.getRefreshToken();
            return new OAuth2AuthorizationCodeAuthenticationToken(registration, exchange, accessToken, refreshToken, accessTokenResponse.getAdditionalParameters());
        };
    }
}
